import pandas as pd
import chinese_calendar

import time
import os
from typing import Tuple
# import traceback

from tdx_sdk import *
from rules import *


class BackTester():
    '''回测工具类\n
        初始参数: \n
        stock-股票, start-起始时间, end-结束时间, period-收盘价取样周期, ma_in-买入MA值，ma_out-卖出MA值，rule - 策略规则序号\n
        使用方法:\n
        prepareData()\n
        run()\n
    '''

    def __init__(self, stock='', start='', end='', period=30, ma_in=5, ma_out=5, rule=1, log=None):
        self.stock, self.start, self.end, self.period = stock, start, end, period
        self.ma_in, self.ma_out = ma_in, ma_out

        self.rule = rule
        self.log_f = log

        self.data = pd.DataFrame(columns=['price', 'time'])  # 数据集
        self.data_idx = 0   # 数据集中开始运行回测的起点序号
        self.output = pd.DataFrame(columns=[
            'time_in', 'price_in', 'time_out', 'price_out', 'rate_of_return', 'evaluate'])  # 输出的交易记录
        self.count = 0  # 总交易次数
        self.count_s = 0  # 成功次数
        self.count_f = 0  # 失败次数

    def prepareData(self) -> bool:
        """准备执行回测规格的数据集\n
            参数：s-股票编码，start-开始日期，end-结束日期，ma-MA值，p-周期
            返回：[数据集，回测数据起点序号]
            """
        s, start, end, ma_in, ma_out, p = self.stock, self.start, self.end, self.ma_in, self.ma_out, self.period

        data_pathfile = '.tmp/data_%s_%s_%s_%d_%d_%d.csv' % (s, start, end, p, ma_in, ma_out)
        if os.path.exists(data_pathfile):
            self.data = pd.read_csv(data_pathfile,index_col=0)
            return True

        ma = max(ma_in, ma_out)
        # 存入用于计算MA的前序日期
        days = pd.DatetimeIndex(pd.date_range(
            end=start, periods=int(ma*p/240)+2, freq='B', closed='left'))
        # 存入要回测的数据日期区间
        days = days.append(pd.DatetimeIndex(
            pd.date_range(start, end, freq='B')))

        self.log_f("Fetch the stock data: " +
                   '{:%Y%m%d}'.format(days[0]) + '-' + end)

        for di in days:
            if di == pd.to_datetime(start):
                self.data_idx = len(self.data)

            if chinese_calendar.is_holiday(di):  # 跳过节假日
                continue

            try:
                rd = pd.DataFrame()

                # 优先读取本地缓存数据
                date = di.strftime('%Y%m%d')
                path = '.data/minutes/%s/' % s
                pathfile = path + '%s.csv' % date
                if os.path.exists(pathfile):
                    rd = pd.read_csv(pathfile)
                else:  # 本地无缓存，调用服务器接口查询
                    rd = tdx_client.minutes(s, date)  # 每天每分钟的交易价格
                    time.sleep(CLIENT_QUERY_PERIOD)  # 控制一下服务器访问频率
                    os.path.exists(path) or os.makedirs(path)
                    rd.to_csv(pathfile)

                if rd is not None:
                    for i in range(1, int(240/p+1)):
                        idx = i*p-1
                        self.data = self.data.append(
                            {'price': rd.at[idx, 'price'], 'time': _idx2time(di, idx)}, ignore_index=True)
            except:
                # traceback.print_exc()
                self.log_f("没有查到股票（%s）在日期（%s）的历史交易记录！" %
                           (s, di.strftime('%Y%m%d')))

        self.data.to_csv('.tmp/data_%s_%s_%s_%d_%d_%d.csv' %
                         (s, start, end, p, ma_in, ma_out))
        if len(self.data) == 0:
            return False
        else:
            return True

    def execTransactions(self) -> bool:
        """回测：按照规则执行交易
        参数：s-股票代码，start-回测开始时间，end-回测结束时间，p-周期（30、60、240），ma-MA值
        返回：True: 有交易产生
        """

        # 逐条记录执行规则rule_f
        rule_f = globals()['rule_%d' % self.rule]
        init_context_rule()
        i = self.data_idx
        while i < len(self.data):
            result, self.output = rule_f(i, self.data, self.output,  self.ma_in, self.ma_out)
            if result < 0:
                self.count, self.count_f = (self.count+1, self.count_f + 1)
            else:
                self.count, self.count_s = (self.count+result, self.count_s+result)
            i = i + 1

        # 输出结果到csv文件
        self.output.to_csv('.tmp/output_%s_%s_%s_%d_%d_%d_%d.csv' % (
            self.stock, self.start, self.end, self.period, self.ma_in, self.ma_out, self.rule))
        if self.count > 0:
            return True
        else:
            return False


def _idx2time(dt, index) -> pd.Timestamp:
    """日分钟序号转为时间
    参数：dt-日期，index-分钟序号，0~239，每天从9:31~15:00
    返回：YY-MM-DD HH：MM
    """
    if index < 29:
        return dt.replace(hour=9, minute=31+index)
    elif index < 89:
        return dt.replace(hour=10, minute=index-29)
    elif index < 120:
        return dt.replace(hour=11, minute=index-89)
    elif index < 179:
        return dt.replace(hour=13, minute=index-119)
    elif index < 239:
        return dt.replace(hour=14, minute=index-179)
    else:
        return dt.replace(hour=15, minute=index-239)

# 测试


def _test():
    # 159995, 20200210上市
    # b = BackTester('159995', '20211021', '20211026', 30, 5, 1, print)
    b = BackTester('159995', '20210101', '20211026', 120, 10, 5, 1, print)

    # 512880, 20161008上市
    # b = BackTester( '512880', '20180801', '20211026',30,5,10,1,print)

    time_start = time.time()
    b.prepareData() and b.execTransactions()
    time_end = time.time()
    print("运行耗费时长（s）：%d" % (time_end - time_start))

    print('Result of backtesting (%s, %s-%s, %dm, MA%d, MA%d):' %
          (b.stock, b.start, b.end, b.period, b.ma_in, b.ma_out))
    print("总交易次数：%d 成功次数：%d 失败次数：%d" %
          (b.count, b.count_s, b.count_f))
    if b.count != 0:
        print("成功率：%.2f%%" % round(float(b.count_s)/b.count*100, 2))
    else:
        print("成功率：NA")


if __name__ == "__main__":
    _test()
